package _12_avl_tree.case2;

/**
 * @author: mornd
 * @dateTime: 2023/6/15 - 15:08
 */
public class AVLTree<E> {
    public AVLNode root;

    /**
     * 新增节点
     *
     * @param k
     * @param v
     */
    public void put(int k, E v) {
        root = put(root, k, v);
    }

    private AVLNode put(AVLNode n, int k, E v) {
        // 1、找到空位，直接插入
        if (n == null) {
            return new AVLNode(k, v);
        }

        if (k < n.key) {
            n.left = put(n.left, k, v);
        } else if (k > n.key) {
            n.right = put(n.right, k, v);
        } else {
            // key重复，覆盖value
            n.value = v;
            return n;
        }

        updateHeight(n); // 更新高度
        return balance(n); // 将自平衡后的节点返回
    }

    /**
     * 删除节点
     * 参考图片：img\删除前.png
     * @param k
     */
    public void remove(int k) {
        root = remove(root, k);
    }

    private AVLNode remove(AVLNode node, int k) {
        // 没有找到节点
        if(node == null) {
            return null;
        }
        if(k < node.key) {
            node.left = remove(node.left, k);
        } else if(k > node.key) {
            node.right = remove(node.right, k);
        } else {
            // 找到了
            if(node.left == null && node.right == null) {
                return null;
            } else if(node.left == null) {
                node = node.right;
            } else if(node.right == null) {
                node = node.left;
            } else {
                // 待删除节点的左右孩子都有
                // 找后继节点
                AVLNode s = node.right;
                while (s.left != null) {
                    s = s.left;
                }
                // 如果后继节点就是 node.right，则返回 node.right 的子节点
                s.right = remove(node.right, s.key);
                s.left = node.left;
                node = s;
            }
         }

        updateHeight(node);
        return balance(node);
    }

    /**
     * 查找节点值
     */
    public E get(int key) {
        AVLNode curr = root;
        while (curr != null) {
            if(key < curr.key) {
                curr = curr.left;
            } else if(key > curr.key) {
                curr = curr.right;
            } else {
                return curr.value;
            }
        }
        return null;
    }

    /**
     * 判断该节点是否失衡，如果失衡则让它保持平衡
     *
     * @param node
     * @return 平衡后的根节点
     */
    private AVLNode balance(AVLNode node) {
        if (node == null) {
            return null;
        }

        // 获取平衡因子
        int bf = balanceFactor(node);
        // 失衡的四种情况处理 LL LR RL RR
        if (bf > 1 && balanceFactor(node.left) >= 0) { // LL
            return rightRotate(node);
        } else if (bf > 1 && balanceFactor(node.left) < 0) { // LR
            return leftRightRotate(node);
        } else if (bf < -1 && balanceFactor(node.right) > 0) { // RL
            return rightLeftRotate(node);
        } else if (bf < -1 && balanceFactor(node.right) <= 0) { // RR
            return leftRotate(node);
        }
        return node;
    }

    /**
     * 左旋
     *
     * @param curr 待旋转的节点
     * @return 新的根节点
     */
    private AVLNode leftRotate(AVLNode curr) {
        AVLNode rightChild = curr.right;
        AVLNode leftChild = rightChild.left; // leftChild 可以为 null
        rightChild.left = curr; // rightChild 上位替代 curr
        curr.right = leftChild; // leftChild 换爹为 curr

        // 更新高度，只有这两个节点的高度需要更新
        updateHeight(curr);
        updateHeight(rightChild);
        return rightChild;
    }

    /**
     * 右旋
     */
    private AVLNode rightRotate(AVLNode curr) {
        AVLNode leftChild = curr.left;
        AVLNode rightChild = leftChild.right;
        leftChild.right = curr;
        curr.left = rightChild;

        updateHeight(curr);
        updateHeight(leftChild);
        return leftChild;
    }

    /**
     * 先左旋左子树，再右旋根节点
     *
     * @param curr
     * @return
     */
    private AVLNode leftRightRotate(AVLNode curr) {
        curr.left = leftRotate(curr.left);
        return rightRotate(curr);
    }

    /**
     * 先右旋右子树，再左旋根节点
     *
     * @param curr
     * @return
     */
    private AVLNode rightLeftRotate(AVLNode curr) {
        curr.right = rightRotate(curr.right);
        return leftRotate(curr);
    }

    /**
     * 获取节点的平衡因子      左子节点 - 右子节点 = 平衡因子
     * 平衡因子为 1，0，-1 表示已经平衡，否则不平衡
     * >= 1表示左边偏高，等于0表示左右一样高，<= -1表示右边偏高
     *
     * @param node
     * @return
     */
    private int balanceFactor(AVLNode node) {
        return height(node.left) - height(node.right);
    }

    /**
     * 更新节点高度
     */
    public void updateHeight(AVLNode node) {
        node.height = 1 + Math.max(height(node.left), height(node.right));
    }

    /**
     * 求节点高度
     */
    public int height(AVLNode node) {
        return node == null ? 0 : node.height;
    }

    /**
     * 树节点
     */
    public class AVLNode {
        public int key;
        public E value;
        public AVLNode left;
        public AVLNode right;
        public int height = 1;

        public AVLNode(int key, E value) {
            this.key = key;
            this.value = value;
        }

        public AVLNode(int key) {
            this.key = key;
        }

        public AVLNode(int key, int height) {
            this.key = key;
            this.height = height;
        }
    }

//    @Override
//    public String toString() {
//        StringBuilder sb = new StringBuilder();
//        traverse(root, "", "", sb);
//        return sb.toString();
//    }
//
//    // 中序遍历打印树结构
//    private void traverse(AVLNode node, String prefix, String childrenPrefix, StringBuilder sb) {
//        if (node != null) {
//            traverse(node.right, childrenPrefix + "├── ", childrenPrefix + "│   ", sb);
//            sb.append(prefix).append(node.key).append("\n");
//            traverse(node.left, childrenPrefix + "└── ", childrenPrefix + "    ", sb);
//        }
//    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        printVerticalTree(root, "", sb);
        return sb.toString();
    }

    private void printVerticalTree(AVLNode node, String prefix, StringBuilder sb) {
        if (node == null) {
            return;
        }

        printVerticalTree(node.right, prefix + "        ", sb);
        sb.append(prefix).append("|--").append(node.key).append("\n");
        printVerticalTree(node.left, prefix + "        ", sb);
    }
}
